package com.roy.sparkDemos.ml;

import org.apache.spark.SparkConf;
import org.apache.spark.api.java.JavaRDD;
import org.apache.spark.api.java.JavaSparkContext;
import org.apache.spark.ml.evaluation.RegressionEvaluator;
import org.apache.spark.ml.recommendation.ALS;
import org.apache.spark.ml.recommendation.ALSModel;
import org.apache.spark.sql.Dataset;
import org.apache.spark.sql.Row;
import org.apache.spark.sql.SparkSession;

public class JavaAlSExample {

    /**
     * ./spark-submit --master yarn --class "com.roy.sparkDemos.ml.JavaAlSExample"
     */

    public static void main(String[] args){
        SparkConf sparkConf = new SparkConf();
        sparkConf.setMaster("local[4]");
        sparkConf.setAppName("MySPWordCount");

        JavaSparkContext sparkContext = new JavaSparkContext(sparkConf);

        SparkSession sparkSession = SparkSession.builder().config(sparkConf).getOrCreate();

        JavaRDD<Object> javaRDD = sparkContext.textFile("hdfs://master:9000/roy/als/sample_movielens_ratings.txt")
                .map(line -> {
                    return Rating.parseLine(line);
                });
        Dataset<Row> dataFrame = sparkSession.createDataFrame(javaRDD, Rating.class);
        Dataset<Row>[] dataSplits = dataFrame.randomSplit(new double[]{0.8, 0.2});
        Dataset<Row> trainingRDD = dataSplits[0];
        Dataset<Row> testRDD = dataSplits[1];

        ALS als =new ALS();
        als.setMaxIter(5).setRegParam(0.01)
                .setUserCol("userId")
                .setItemCol("prodId")
                .setRatingCol("rating");
        ALSModel alsModel = als.fit(trainingRDD);

        alsModel.setColdStartStrategy("drop");
        Dataset<Row> predictions = alsModel.transform(testRDD);

        RegressionEvaluator evaluator = new RegressionEvaluator()
                .setMetricName("rmse")
                .setLabelCol("rating")
                .setPredictionCol("prediction");
        Double rmse = evaluator.evaluate(predictions);
        System.out.println("Root-mean-square error = " + rmse);

        // Generate top 10 movie recommendations for each user
        Dataset<Row> userRecs = alsModel.recommendForAllUsers(10);
        // Generate top 10 user recommendations for each movie
        Dataset<Row> movieRecs = alsModel.recommendForAllItems(10);
        // $example off$
        userRecs.show();
        movieRecs.show();

        sparkSession.stop();


    }
}
